"""
https://keras.io/zh/examples/cifar10_cnn/
在 CIFAR10 小型图像数据集上训练一个深度卷积神经网络。
"""
import os
import keras
from keras.datasets import cifar10
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.utils import to_categorical
from keras.layers import (Dense, Dropout, Activation, Flatten,
                          Conv2D, MaxPooling2D)
from utils.plot_graph import keras_plot_model
from utils.search_model_paramter import grid_search
from utils.model_fit import train_model


def load_data(num_classes):
    """

    :param num_classes:
    :return:
    """
    (x_train, y_train), (x_test, y_test) = cifar10.load_data()

    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')
    x_train /= 255
    x_test /= 255

    y_train = to_categorical(y_train, num_classes=num_classes)
    y_test  = to_categorical(y_test , num_classes=num_classes)

    return x_train, y_train, x_test, y_test


def build_fit_model(x_train, y_train, x_test, y_test,num_classes,
                    unit, act, opt, loss, epochs, batch_size, verbose,
                    pic_path, model_name, workers=4,
                    use_multiprocessing=True, shuffle=True, plot=False,
                    **options):
    """

    :param train_x:
    :param train_y:
    :param test_x:
    :param test_y:
    :param num_classes:
    :param unit:
    :param act:
    :param opt:
    :param loss:
    :param epochs:
    :param batch_size:
    :param verbose:
    :param pic_path:
    :param model_name:
    :param kernel_size:
    :param workers:
    :param use_multiprocessing:
    :param shuffle:
    :param plot:
    :return:
    """
    model = Sequential()
    model.add(Conv2D(filters=32,
                     kernel_size=(3, 3),
                     padding='same',
                     activation='relu',
                     input_shape=(x_train.shape[1:])))
    model.add(Conv2D(filters=32,
                     kernel_size=(3, 3),
                     activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))

    model.add(Conv2D(filters=64,
                     kernel_size=(3, 3),
                     padding='same',
                     activation='relu',
                     input_shape=(x_train.shape[1:])))
    model.add(Conv2D(filters=64,
                     kernel_size=(3, 3),
                     activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))

    model.add(Flatten())
    model.add(Dense(units=512,
                    activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(units=num_classes,
                    activation='softmax'))

    model.compile(loss='categorical_crossentropy',
                  optimizer='rmsprop',
                  metrics=['accuracy'])

    return train_model(model=model,
                       train_x=x_train, train_y=y_train,
                       test_x=x_test, test_y=y_test,
                       epochs=epochs,
                       model_name=model_name,
                       pic_path=pic_path,
                       shuffle=shuffle,
                       verbose=verbose,
                       plot=plot,
                       batch_size=batch_size,
                       workers=workers,
                       use_multiprocessing=use_multiprocessing)


def conv_cifar_run(workers):
    """
    主程序
    :param workers:
    :return:
    """
    get_parameter = False
    train_model = True

    kernel_size_arr = [[(2, 2), (2, 2, 2)],
                       [(3, 3), (3, 3, 3)],
                       [(4, 4), (4, 4, 4)]]
    unit_arr = [30, 40, 50]
    act_arr = ['sigmoid']
    opt_arr = ['RMSprop', 'Adam', 'Adadelta', 'Nadam']
    loss_arr = ['binary_crossentropy']
    """
    padding_arr = ['valid', 'same']

    valid：不填充,尺寸会变小
    same：进行左右上下的补齐, 输入和输出尺寸会保持一致
        其中左,上依次补齐 flood (kernel_size -1 ) / 2 , 
        右下补齐ceil (( kernel_size - 1) /2 ) ,
        补齐后进行的操作就是类似valid下的滑动卷积
        output_shape = ceil (input_shape / stride)
    """
    epochs = 100
    batch_size = 32
    num_classes = 10

    x_train, y_train, x_test, y_test = load_data(num_classes=num_classes)

    # 参数选择
    if get_parameter:
        for kernel_size in kernel_size_arr:
            print('-' * 100)
            print('kernel_size: {}'.format(kernel_size))
            grid_search(func=build_fit_model,
                        train_x=x_train, train_y=y_train,
                        test_x=x_test, test_y=y_test,
                        unit_arr=unit_arr,
                        act_arr=act_arr,
                        opt_arr=opt_arr,
                        loss_arr=loss_arr,
                        epochs=epochs,
                        batch_size=batch_size,
                        workers=workers,
                        use_multiprocessing=True,
                        shuffle=True,
                        plot=False,
                        kernel_size=kernel_size)

    # 模型训练
    if train_model:
        model, loss_metrics = \
            build_fit_model(x_train=x_train, y_train=y_train,
                            x_test=x_test, y_test=y_test,
                            num_classes=num_classes,
                            unit=60,
                            act='sigmoid',
                            opt='RMSprop',
                            loss='binary_crossentropy',
                            pic_path='../pic',
                            model_name='CNN_LSTM',
                            epochs=epochs,
                            batch_size=batch_size,
                            verbose=1,
                            workers=workers,
                            use_multiprocessing=True,
                            shuffle=True,
                            plot=True,
                            kernel_size=[(4, 4), (4, 4, 4)])
        print('loss_metrics: {}'.format(loss_metrics))
        # plot model
        keras_plot_model(model=model)
        # keras_plot_model(model=model, pic='../pic/CNN_LSTM.png')
